"""
@description: Monitor the process of training, and record related indicators.
@author: Zzay
@create: 2022/07/09
"""
from tensorflow.keras.callbacks import BaseLogger
import numpy as np
import matplotlib.pyplot as plt
import json
import os


class TrainingMonitor(BaseLogger):
    """
    Class that monitor the process of training, and do some records.
    """

    def __init__(self, figure_path, json_path=None, start_at=0):
        """
        Store the output path for the figure, the JSON serialized file, and the starting epoch.
        :param figure_path: Path for the figure outputted
        :param json_path: Path for the JSON serialized file
        :param start_at: Epoch where monitor starts
        """
        super(TrainingMonitor, self).__init__()
        self.figure_path = figure_path
        self.json_path = json_path
        self.start_at = start_at

    def on_train_begin(self, logs={}):
        """
        Called once the training process starts.
        :param logs: Logs
        :return: Record logs
        """
        # initialize the history dictionary
        self.history = {}

        # if the JSON history path exists, load the training history
        if self.json_path is not None:
            if os.path.exists(self.json_path):
                self.history = json.loads(open(self.json_path).read())
                # check to see if a starting epoch was supplied
                if self.start_at > 0:
                    # loop over the entries in the history log,
                    # and trim any entries that are past the starting epoch
                    for i in self.history.keys():
                        self.history[i] = self.history[i][:self.start_at]

    def on_epoch_end(self, epoch, logs={}):
        """
        Called when each epoch ends.
        :param epoch: Current epoch order
        :param logs: Logs
        :return: Record logs
        """
        # loop over the logs, and update the loss, accuracy, etc. for the entire training process
        for (k, v) in logs.items():
            loss = self.history.get(k, [])
            loss.append(v)
            self.history[k] = loss

        # check to see if the training history should be serialized to file
        if self.json_path is not None:
            file = open(self.json_path, "w")
            file.write(json.dumps(self.history))
            file.close()

        # ensure at least two epochs have passed before plotting (starts at 0)
        # plot the training loss and accuracy
        N = np.arange(1, len(self.history["loss"]) + 1)
        plt.style.use("ggplot")
        file_name = os.path.splitext(self.figure_path)[0]
        file_extent = os.path.splitext(self.figure_path)[1]
        loss_plot_file = file_name + '_loss' + file_extent
        accuracy_plot_file = file_name + '_accuracy' + file_extent

        # loss curve
        plt.figure()
        plt.plot(N, self.history["loss"], label="train_loss")
        plt.plot(N, self.history["val_loss"], label="val_loss")
        plt.title("Training Loss [Epoch {}]".format(len(self.history["loss"])))
        plt.xlabel("Epoch #")
        plt.ylabel("Loss")
        plt.legend()
        # save the figure
        plt.savefig(loss_plot_file)
        plt.close()

        # accuracy curve
        plt.figure()
        plt.plot(N, self.history["accuracy"], label="train_acc")
        plt.plot(N, self.history["val_accuracy"], label="val_acc")
        plt.title("Training Accuracy [Epoch {}]".format(len(self.history["loss"])))
        plt.xlabel("Epoch #")
        plt.ylabel("Accuracy")
        plt.legend()
        # save the figure
        plt.savefig(accuracy_plot_file)
        plt.close()
